"""
会议调度Agent环境

模拟真实的会议安排场景，支持：
- 团队成员日程查询
- 历史会议模式分析
- 协商策略执行
- 会议创建和验证
"""

import json
import random
import numpy as np
from typing import Dict, List, Tuple, Optional
from datetime import datetime, timedelta


class MeetingSchedulingEnv:
    """
    会议调度环境

    状态空间：
    - 会议主题
    - 时间范围（一周，15个时间段：5天x3时段）
    - 必选/可选参会人员
    - 日程可用性矩阵
    - 历史会议模式
    - 协商历史

    动作空间：
    - query_calendar: 查询人员日程
    - query_history: 查询历史会议
    - negotiate_time: 协商时间段
    - select_time: 选择会议时间
    - create_meeting: 创建会议
    """

    DAYS = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday']
    TIME_SLOTS = ['morning', 'afternoon', 'evening']

    # 人员姓名池
    NAME_POOL = ['李明', '王芳', '陈静', '刘强', '赵敏', '张伟', '周娟', '吴杰']

    # 会议主题池
    TOPICS = [
        '产品规划会', '技术评审会', '季度总结会', '项目启动会',
        '需求讨论会', '设计评审会', '运营周会', '客户沟通会'
    ]

    def __init__(self,
                 max_steps: int = 20,
                 num_total_people: int = 6,
                 seed: Optional[int] = None):
        """
        参数:
            max_steps: 最大步数
            num_total_people: 团队总人数
            seed: 随机种子
        """
        self.max_steps = max_steps
        self.num_total_people = num_total_people
        self.seed_value = seed

        if seed is not None:
            random.seed(seed)
            np.random.seed(seed)

        # 环境状态
        self.current_step = 0
        self.meeting_topic = None
        self.required_attendees = []
        self.optional_attendees = []
        self.all_people = []

        # 日程矩阵: {person: {day: {time_slot: status}}}
        self.calendar_data = {}

        # 历史会议数据
        self.historical_meetings = []
        self.historical_pattern = None  # (day, time_slot)

        # 运行时状态
        self.negotiation_history = []
        self.queried_calendars = set()
        self.queried_history = False
        self.selected_time = None
        self.meeting_created = False

        # 观察缓存
        self.last_observation = ""

    def reset(self, instance: Optional[Dict] = None) -> Tuple[str, Dict]:
        """
        重置环境

        参数:
            instance: 可选的实例配置，包含：
                - meeting_topic: 会议主题
                - required_attendees: 必选人员列表
                - optional_attendees: 可选人员列表
                - calendar_data: 日程数据
                - historical_meetings: 历史会议记录

        返回:
            (observation, info): 初始观察和信息
        """
        self.current_step = 0
        self.negotiation_history = []
        self.queried_calendars = set()
        self.queried_history = False
        self.selected_time = None
        self.meeting_created = False

        if instance is None:
            instance = self._generate_random_instance()

        # 加载实例数据
        self.meeting_topic = instance['meeting_topic']
        self.required_attendees = instance['required_attendees']
        self.optional_attendees = instance['optional_attendees']
        self.all_people = self.required_attendees + self.optional_attendees
        self.calendar_data = instance['calendar_data']
        self.historical_meetings = instance.get('historical_meetings', [])

        # 计算历史模式
        self._compute_historical_pattern()

        # 初始观察
        observation = self._build_initial_observation()
        self.last_observation = observation

        info = {
            'meeting_topic': self.meeting_topic,
            'required_count': len(self.required_attendees),
            'optional_count': len(self.optional_attendees),
            'historical_pattern': self.historical_pattern,
            'step': self.current_step
        }

        return observation, info

    def _generate_random_instance(self) -> Dict:
        """生成随机会议安排实例"""
        # 随机选择会议主题
        topic = random.choice(self.TOPICS)

        # 随机选择人员
        selected_people = random.sample(self.NAME_POOL,
                                       min(self.num_total_people, len(self.NAME_POOL)))

        # 随机分配必选/可选
        num_required = random.randint(2, min(4, len(selected_people) - 1))
        required = selected_people[:num_required]
        optional = selected_people[num_required:]

        # 生成日程数据
        calendar_data = {}
        for person in selected_people:
            calendar_data[person] = {}
            for day in self.DAYS:
                calendar_data[person][day] = {}
                for slot in self.TIME_SLOTS:
                    # 60%概率空闲
                    status = 'available' if random.random() < 0.6 else 'busy'
                    calendar_data[person][day][slot] = status

        # 生成历史会议（过去5周）
        historical_meetings = []
        historical_day_slot = random.choice([(d, s) for d in self.DAYS for s in self.TIME_SLOTS])

        for week in range(-5, 0):
            # 80%概率在历史偏好时段
            if random.random() < 0.8:
                day, slot = historical_day_slot
            else:
                day = random.choice(self.DAYS)
                slot = random.choice(self.TIME_SLOTS)

            historical_meetings.append({
                'topic': topic,
                'week': week,
                'day': day,
                'time_slot': slot
            })

        return {
            'meeting_topic': topic,
            'required_attendees': required,
            'optional_attendees': optional,
            'calendar_data': calendar_data,
            'historical_meetings': historical_meetings
        }

    def _compute_historical_pattern(self):
        """计算历史会议模式"""
        if not self.historical_meetings:
            self.historical_pattern = None
            return

        # 统计每个(day, time_slot)的频次
        pattern_counts = {}
        for meeting in self.historical_meetings:
            key = (meeting['day'], meeting['time_slot'])
            pattern_counts[key] = pattern_counts.get(key, 0) + 1

        # 找出最高频次
        if pattern_counts:
            self.historical_pattern = max(pattern_counts.items(), key=lambda x: x[1])[0]
        else:
            self.historical_pattern = None

    def _build_initial_observation(self) -> str:
        """构建初始观察（用户请求）"""
        required_str = '、'.join(self.required_attendees)
        optional_str = '、'.join(self.optional_attendees)

        observation = (
            f"用户请求：下周找个时间开{self.meeting_topic}，"
            f"必须参加的有{required_str}，"
            f"可选参加的有{optional_str}。"
        )

        return observation

    def step(self, action_dict: Dict, valid: int) -> Tuple[str, float, bool, Dict]:
        """
        执行一步动作

        参数:
            action_dict: 动作字典，格式如：
                {
                    "action_type": "query_calendar" | "query_history" | "negotiate_time" | "select_time",
                    "parameters": {...}
                }
            valid: 动作是否有效 (0或1)

        返回:
            (observation, reward, done, info)
        """
        self.current_step += 1

        # 无效动作直接返回
        if not valid:
            observation = "错误：动作格式无效，请按照规定的JSON格式提供动作。"
            return observation, -1.0, True, {'error': 'invalid_action', 'step': self.current_step}

        # 超过最大步数
        if self.current_step > self.max_steps:
            observation = "错误：已超过最大步数限制。"
            return observation, -10.0, True, {'error': 'max_steps', 'step': self.current_step}

        action_type = action_dict.get('action_type', '')
        parameters = action_dict.get('parameters', {})

        # 路由到对应的动作处理函数
        if action_type == 'query_calendar':
            return self._action_query_calendar(parameters)
        elif action_type == 'query_history':
            return self._action_query_history(parameters)
        elif action_type == 'negotiate_time':
            return self._action_negotiate_time(parameters)
        elif action_type == 'select_time':
            return self._action_select_time(parameters)
        else:
            observation = f"错误：未知的动作类型 '{action_type}'。"
            return observation, -1.0, True, {'error': 'unknown_action', 'step': self.current_step}

    def _action_query_calendar(self, parameters: Dict) -> Tuple[str, float, bool, Dict]:
        """查询日程动作"""
        persons = parameters.get('persons', [])

        if not persons:
            return "错误：未指定要查询的人员。", -0.5, False, {'step': self.current_step}

        # 查询每个人的日程
        calendar_results = {}
        for person in persons:
            if person not in self.all_people:
                continue

            self.queried_calendars.add(person)
            calendar_results[person] = self.calendar_data[person]

        # 格式化输出
        observation = "日程查询结果：\n"
        for person, schedule in calendar_results.items():
            observation += f"\n{person}的日程：\n"
            for day in self.DAYS:
                day_schedule = []
                for slot in self.TIME_SLOTS:
                    status = schedule[day][slot]
                    day_schedule.append(f"{slot}:{status}")
                observation += f"  {day}: {', '.join(day_schedule)}\n"

        info = {
            'step': self.current_step,
            'action': 'query_calendar',
            'persons': persons
        }

        return observation, 0.0, False, info

    def _action_query_history(self, parameters: Dict) -> Tuple[str, float, bool, Dict]:
        """查询历史会议动作"""
        if self.queried_history:
            return "提示：历史会议已查询过。", 0.0, False, {'step': self.current_step}

        self.queried_history = True

        # 统计历史模式
        if not self.historical_meetings:
            observation = "历史会议查询结果：暂无历史记录。"
        else:
            pattern_counts = {}
            for meeting in self.historical_meetings:
                key = (meeting['day'], meeting['time_slot'])
                pattern_counts[key] = pattern_counts.get(key, 0) + 1

            observation = f"历史会议查询结果（{self.meeting_topic}）：\n"
            for meeting in self.historical_meetings:
                observation += f"  第{meeting['week']}周: {meeting['day']}, {meeting['time_slot']}\n"

            if self.historical_pattern:
                day, slot = self.historical_pattern
                count = pattern_counts[self.historical_pattern]
                percentage = (count / len(self.historical_meetings)) * 100
                observation += f"\n历史规律：{percentage:.0f}%的{self.meeting_topic}安排在{day}的{slot}。"

        info = {
            'step': self.current_step,
            'action': 'query_history',
            'historical_pattern': self.historical_pattern
        }

        return observation, 0.0, False, info

    def _action_negotiate_time(self, parameters: Dict) -> Tuple[str, float, bool, Dict]:
        """协商时间动作"""
        person = parameters.get('person')
        day = parameters.get('day')
        time_slot = parameters.get('time_slot')

        if not all([person, day, time_slot]):
            return "错误：协商参数不完整（需要person, day, time_slot）。", -0.5, False, {'step': self.current_step}

        if person not in self.all_people:
            return f"错误：人员'{person}'不在参会人员列表中。", -0.5, False, {'step': self.current_step}

        if day not in self.DAYS or time_slot not in self.TIME_SLOTS:
            return f"错误：时间参数无效。", -0.5, False, {'step': self.current_step}

        # 协商成功（简化：总是成功）
        original_status = self.calendar_data[person][day][time_slot]
        self.calendar_data[person][day][time_slot] = 'available'

        self.negotiation_history.append({
            'person': person,
            'day': day,
            'time_slot': time_slot,
            'original_status': original_status,
            'result': 'success'
        })

        observation = f"协商结果：成功与{person}协商{day}的{time_slot}时段，已预留时间。"

        info = {
            'step': self.current_step,
            'action': 'negotiate_time',
            'person': person,
            'negotiations_count': len(self.negotiation_history)
        }

        return observation, 0.0, False, info

    def _action_select_time(self, parameters: Dict) -> Tuple[str, float, bool, Dict]:
        """选择时间并创建会议"""
        day = parameters.get('day')
        time_slot = parameters.get('time_slot')

        if not all([day, time_slot]):
            return "错误：时间参数不完整（需要day, time_slot）。", -1.0, True, {'step': self.current_step}

        if day not in self.DAYS or time_slot not in self.TIME_SLOTS:
            return f"错误：时间参数无效。", -1.0, True, {'step': self.current_step}

        self.selected_time = (day, time_slot)

        # 统计参会情况
        required_available = []
        required_busy = []
        for person in self.required_attendees:
            if self.calendar_data[person][day][time_slot] == 'available':
                required_available.append(person)
            else:
                required_busy.append(person)

        optional_available = []
        optional_busy = []
        for person in self.optional_attendees:
            if self.calendar_data[person][day][time_slot] == 'available':
                optional_available.append(person)
            else:
                optional_busy.append(person)

        # 计算奖励
        reward = self._calculate_reward(
            required_available=len(required_available),
            required_total=len(self.required_attendees),
            optional_available=len(optional_available),
            optional_total=len(self.optional_attendees),
            selected_day=day,
            selected_slot=time_slot,
            negotiations_count=len(self.negotiation_history)
        )

        # 构建观察
        all_required_met = len(required_busy) == 0

        if all_required_met:
            observation = (
                f"✓ 会议创建成功！\n"
                f"时间：{day}, {time_slot}\n"
                f"必选参会人员({len(required_available)}/{len(self.required_attendees)}): "
                f"{', '.join(required_available)}\n"
                f"可选参会人员({len(optional_available)}/{len(self.optional_attendees)}): "
                f"{', '.join(optional_available) if optional_available else '无'}\n"
                f"协商次数：{len(self.negotiation_history)}\n"
                f"奖励得分：{reward:.1f}"
            )
            self.meeting_created = True
        else:
            observation = (
                f"✗ 会议创建失败！\n"
                f"时间：{day}, {time_slot}\n"
                f"必选人员中有{len(required_busy)}人不可用：{', '.join(required_busy)}\n"
                f"奖励得分：{reward:.1f}"
            )

        # 检查历史匹配
        historical_match = (
            self.historical_pattern is not None and
            self.historical_pattern == (day, time_slot)
        )

        info = {
            'step': self.current_step,
            'action': 'select_time',
            'selected_time': (day, time_slot),
            'required_met': all_required_met,
            'required_available': len(required_available),
            'required_total': len(self.required_attendees),
            'optional_available': len(optional_available),
            'optional_total': len(self.optional_attendees),
            'historical_match': historical_match,
            'negotiations_count': len(self.negotiation_history),
            'reward': reward,
            'won': all_required_met  # 成功标志
        }

        return observation, reward, True, info

    def _calculate_reward(self,
                         required_available: int,
                         required_total: int,
                         optional_available: int,
                         optional_total: int,
                         selected_day: str,
                         selected_slot: str,
                         negotiations_count: int) -> float:
        """
        计算奖励

        奖励组成：
        1. 必选人员参与率（权重最高）：-100 per缺席
        2. 可选人员参与率：+1 per参与
        3. 历史模式匹配：+5
        4. 协商成本：-2 per协商
        """
        reward = 0.0

        # 1. 必选人员（优先级最高）
        if required_available < required_total:
            missing_required = required_total - required_available
            reward -= missing_required * 100

        # 2. 可选人员
        if required_available == required_total:  # 只有在满足必选的前提下才加分
            if optional_total > 0:
                optional_rate = optional_available / optional_total
                reward += optional_rate * optional_total * 1

        # 3. 历史模式
        if required_available == required_total:
            if self.historical_pattern and self.historical_pattern == (selected_day, selected_slot):
                reward += 5

        # 4. 协商成本
        reward -= negotiations_count * 2

        return reward

    def get_state_dict(self) -> Dict:
        """获取完整状态（用于调试）"""
        return {
            'meeting_topic': self.meeting_topic,
            'required_attendees': self.required_attendees,
            'optional_attendees': self.optional_attendees,
            'calendar_data': self.calendar_data,
            'historical_pattern': self.historical_pattern,
            'negotiation_history': self.negotiation_history,
            'queried_calendars': list(self.queried_calendars),
            'queried_history': self.queried_history,
            'selected_time': self.selected_time,
            'current_step': self.current_step
        }


if __name__ == '__main__':
    # 测试环境
    env = MeetingSchedulingEnv(seed=42)
    obs, info = env.reset()

    print("=== 初始观察 ===")
    print(obs)
    print(f"\nInfo: {info}")

    print("\n=== 测试动作：查询日程 ===")
    action1 = {
        'action_type': 'query_calendar',
        'parameters': {
            'persons': info['meeting_topic']  # 会获取到具体人员
        }
    }

    # 获取实际人员
    all_people = env.all_people
    action1['parameters']['persons'] = all_people[:2]

    obs, reward, done, info = env.step(action1, valid=1)
    print(obs)
    print(f"Reward: {reward}, Done: {done}")

    print("\n=== 测试动作：查询历史 ===")
    action2 = {
        'action_type': 'query_history',
        'parameters': {}
    }
    obs, reward, done, info = env.step(action2, valid=1)
    print(obs)
    print(f"Reward: {reward}, Done: {done}")
